import pygame
import math
import sys
import copy
import time

# 取符号函数
def sign(a):
    return (a > 0) - (a < 0)


# 万有引力函数
def gravity(p1, p2):
    m1, x1, y1, v_x1, v_y1 = p1
    m2, x2, y2, v_x2, v_y2 = p2
    if (x1, y1) == (x2, y2):
        return p1, p2
    delta_x = (x1 - x2) * k
    delta_y = (y1 - y2) * k
    # 距离平方
    r2 = delta_x ** 2 + delta_y ** 2
    # 万有引力定律
    F = G * m1 * m2 / r2
    # 星球夹角
    theta = math.acos(delta_x / r2 ** 0.5)
    # x、y 轴引力分量
    fx = abs(F * math.cos(theta)) * sign(delta_x)
    fy = abs(F * math.sin(theta)) * sign(delta_y)
    # x、y 轴加速度，牛顿第二定律
    ax2 = fx / m2
    ay2 = fy / m2
    ax1 = fx / m1
    ay1 = fy / m1
    # 速度变化，vt = v0 + at
    v_x2 += ax2 * t
    v_y2 += ay2 * t
    v_x1 -= ax1 * t
    v_y1 -= ay1 * t
    # 位置变化，st = s0 + vt
    x2 += v_x2 * t / k
    y2 += v_y2 * t / k
    x1 -= v_x2 * t / k
    y1 -= v_y2 * t / k
    p3 = [m1, x1, y1, v_x1, v_y1]
    p4 = [m2, x2, y2, v_x2, v_y2]
    return p3, p4


# 定义颜色
black = 0, 0, 0
white = 255, 255, 255

k = 1  # 距离缩放参数
G = 1e-2  # 万有引力常量
t = 1  # 时间缩放参数

# distance_e = 1.495978707e11  # 日地距离
# distance_j = 7.7833e11  # 日木距离
# distance_huo = 2.2794e11  # 日火距离
# distance_shui = 5.791e10  # 日水距离
# distance_jin = 1.082e11  # 日金距离
# distance_tu = 1.4294e12  # 日土距离
# distance_tian = 2.87099e12  # 日天距离
# distance_hai = 4.504e12  # 日海距离

# p = (m, x, y, v_x, v_y)
p1 = [2000000, 500, 300, 0, 0]
shuixing = (p1[1], p1[2])

p2 = [48, 450, 300, -40, 30]
jinxing = (p2[1], p2[2])

p3 = [60, 700, 300, 10, -11]
earth = (p3[1], p3[2])

p4 = [64, 300, 200, 10, 10]
Mars = (p4[1], p4[2])

p5 = [19, 200, 300, 20, 3]
jupiter = (p5[1], p5[2])

p6 = [57, 100, 400, 10, 20]
tuxing = (p6[1], p6[2])

p7 = [87, 50, 500, 40, 0]
tianwang = (p7[1], p7[2])

pygame.init()  # 初始化
screen = pygame.display.set_mode((1000, 600))  # 创建窗口
font = pygame.font.Font(r'C:\Windows\Fonts\simkai.ttf', 30)  # 显示中文需要字体，否则可略过
text = font.render("引力弹弓", 1, white)
pos = font.render(".", 1, white)
pygame.display.set_caption("引力弹弓模拟")

ls_p = [p1, p2, p3, p4, p5, p6, p7]

while True:  # 主循环
    time.sleep(0.1)
    for event in pygame.event.get():
        if event.type in (pygame.QUIT,):
            sys.exit()  # 按键响应，按键后退出
    screen.fill(black)  # 刷新背景
    for i in range(len(ls_p)):
        for j in range(i+1, len(ls_p)):
            ls_p[i], ls_p[j] = gravity(ls_p[i], ls_p[j])
            planet1 = (ls_p[i][1], ls_p[i][2])
            planet2 = (ls_p[j][1], ls_p[j][2])
            print(planet1)
            screen.blit(pos, planet1)
            screen.blit(pos, planet2)
    # screen.blit(font.render('日地距离%.2f×10^8 km' % (distance_e / 1e11), 1, white), (100, 20))
    print('*'*20)
    pygame.display.update()  # 刷新
